package cn.sunxiang0918.tensorflowtest

import org.springframework.shell.standard.ShellComponent
import org.springframework.shell.standard.ShellMethod
import org.springframework.shell.standard.ShellOption
import org.tensorflow.Graph
import org.tensorflow.Session
import org.tensorflow.Tensor
import java.nio.charset.Charset
import java.nio.file.Files
import java.nio.file.Paths
import java.util.*




/**
 * @author sun
 * @version 1.0
 * @Date 2018/2/12 21:05
 */
@ShellComponent("图片分类")
class ClassifierTest {

    private val graphDef: ByteArray by lazy {
        ClassifierTest::class.java.getResource("/models/tensorflow_inception_graph.pb").readBytes()
    }
    private val labels: List<String> by lazy {
        ClassifierTest::class.java.getResource("/models/imagenet_comp_graph_label_strings.txt").readText(Charset.forName("utf-8")).split("\n")
    }

    @ShellMethod(key = ["c"], value = "图片分类")
    fun classifier(@ShellOption(help = "图片路径") img:String) {
        val imgBytes = Files.readAllBytes(Paths.get(img))

        /*图片预处理*/
        val image:Tensor<Float> = constructAndExecuteGraphToNormalizeImage(imgBytes)

        /*找到符合的label以及概率*/
        val labelProbabilities:FloatArray = executeInceptionGraph(graphDef, image)

        /*找到最符合的label*/
        val bestLabelIdx = maxIndex(labelProbabilities)

        println(labels[bestLabelIdx])
    }

    private fun maxIndex(probabilities: FloatArray): Int {
        /*输入的一维数组中每个元素就是可能的概率, 这里找到最高的概率,并记录他的index*/
        var best = 0
        (1 until probabilities.size)
                .asSequence()
                .filter { probabilities[it] > probabilities[best] }
                .forEach { best = it }
        return best
    }

    /**
     * 求图片相近的标签
     */
    private fun executeInceptionGraph(graphDef: ByteArray?, image: Tensor<Float>): FloatArray {
        Graph().use { g ->      //kotlin的use函数 相当于java中的 try-resources,use中的对象必须实现 AutoCloseable 接口
            /*加载模型*/
            g.importGraphDef(graphDef)
            /*把模型graph放入session运行*/
            Session(g).use { s ->
                /*fetch 可以传入一些tensor来传回运行结果
                * feed 机制可以临时替代图中的任意操作中的 tensor 可以对图中任何操作提交补丁, 直接插入一个 tensor. 也就是占位符,可以当做是输入*/
                /*下面这一句话就是执行pb模型中的节点计算,输入时 image, 输出是output*/
                s.runner().feed("input", image).fetch("output").run()[0].expect(java.lang.Float::class.java).use { result ->
                    /*得到输出后,进行处理*/
                    val rshape = result.shape()     //返回结果的shape 形状   TensorFlow文档中使用了三种记号来方便地描述张量的维度：阶，形状以及维数. 比如一个 [[1,2,3],[4,5,6]] 二维数组 它的shape就是 (2, 3)
                    if (result.numDimensions() != 2 || rshape[0] != 1L) {           //校验是否是1维数组 
                        throw RuntimeException(
                                String.format(
                                        "Expected model to produce a [1 N] shaped tensor where N is the number of labels, instead it produced one with shape %s",
                                        Arrays.toString(rshape)))
                    }
                    val nlabels = rshape[1].toInt()     //获取它的阶. 也就是推荐的个数
                    return result.copyTo(Array(1) { FloatArray(nlabels) })[0]       //整理结果,创建一个以为数组,把result拷贝到这个一位数组中去
                }
            }
        }
    }

    /**
     * 重新构造图片,给图片标准化.  inception5h.zip的模型是 按照224*224 RGB来表示的. 所以输入也需要标准化图片
     */
    private inline fun <reified T> constructAndExecuteGraphToNormalizeImage(imgBytes: ByteArray): Tensor<T> {
        Graph().use({ g ->
            val b = GraphBuilder(g)
            // Some constants specific to the pre-trained model at:
            // https://storage.googleapis.com/download.tensorflow.org/models/inception5h.zip
            //
            // - The model was trained with images scaled to 224x224 pixels.
            // - The colors, represented as R, G, B in 1-byte each were converted to
            //   float using (value - Mean)/Scale.
            val H = 224
            val W = 224
            val mean = 117f
            val scale = 1f

            // Since the graph is being constructed once per execution here, we can use a constant for the
            // input image. If the graph were to be re-used for multiple input images, a placeholder would
            // have been more appropriate.
            /*定义图片输入*/
            val input = b.constant("input", imgBytes)
            
            /*定义tensorflow的graph有向无环图.定义其中的操作. 把图像转换成为float数组  float using (value - Mean)/Scale. */
            val output = b.div(     //除法  和scale 相除
                    b.sub(      //减法  和mean 相减
                            b.resizeBilinear(       //图像变形  images: 4维tensor,形状为 [batch, height, width, channels]   size: 1维 int32类型的 Tensor,包含两个元素:new_height, new_width.    ResizeMethod.BILINEAR: 双线性插值(Bilinear interpolation.)  
                                    b.expandDims(           //增加维度, 等于把图片最前端 增加了一个 make_batch 的维度
                                            b.cast(b.decodeJpeg(input, 3), java.lang.Float::class.java),    //类型转换,把输入的图片数据转换为 Float流
                                            b.constant("make_batch", 0)),       //定义批量操作
                                    b.constant("size", intArrayOf(H, W))),      //定义目标尺寸
                            b.constant("mean", mean)),
                    b.constant("scale", scale))
            
            /*把定义好的graph 放入 session中执行*/
            /*fetch 执行输出*/
            /*expect 转换/检测输出结果类型*/
            Session(g).use({ s -> return s.runner().fetch(output.op().name()).run()[0].expect(T::class.java) })
        })
    }

}